#!/usr/bin/env groovy

import groovy.transform.Field

// Configs for build from pytorch docker images
// (See: https://hub.docker.com/r/pytorch/pytorch/tags)
def ubuntu_from_kaolin_configs = [
    [
        'imageTag': 'master-2670-custom-torch1.11.0-cuda11.3.1-cudnn8-py3.8',
        'archsToTest': 'GA100-E4720-DVT'
    ],
]

dockerRegistryServer = 'gitlab-master.nvidia.com:5005'
kaolinDockerRegistryName = 'toronto_dl_lab/kaolin'
kaolinDockerRegistryFull = "${dockerRegistryServer}/${kaolinDockerRegistryName}/kaolin"
wispDockerRegistryName = 'toronto_dl_lab/kaolin-wisp'
wispDockerRegistryFull = "${dockerRegistryServer}/${wispDockerRegistryName}/wisp"

// Used for target docker image tag, as it doesn't support all characters (such as /)
branchRef = gitlabSourceBranch.replaceAll("[^a-zA-Z0-9]", "-")

node {
    checkout scm
    // Sanity check, in case this script fail to launch all builds and tests
    // Right now we only apply CI on MR and master branch.
    // To enable master branch we have to accept all the push requests
    // and prune them here.
    sh "echo ${gitlabActionType}"
    if (gitlabActionType == "MERGE" || gitlabSourceBranch == "main") {
        gitlabCommitStatus("launch all builds") {
            jobMap = [:]
            // Jenkins doesn't parse the commit hash from the webhook.
            // So we need to get the commit hash from the last commit in the branch,
            // So we unsure that all the build and run are on the same commit
            //
            // Note:
            //   If two commits on the same branch are pushed before the first
            //   run this line then they will both run on the second commit
            commitHash = sh(script: "git log -1 --pretty=format:%h",
                            returnStdout: true).trim()
            sh "echo ${commitHash}"
            if (gitlabActionType == "MERGE") {
                sh "echo ${gitlabMergeRequestTitle}"
            }

            for (config in ubuntu_from_kaolin_configs) {
                def configName = "${config['imageTag']}"
                def baseImageTag =  "${kaolinDockerRegistryFull}:${config['imageTag']}"

                jobMap["${configName}"] = prepareUbuntuFromBaseImageJob(
                    configName,
                    baseImageTag,
                    config['archsToTest']
                )
            }

            stage('Launch builds') {
                parallel jobMap
            }
        }
    }
}

def prepareUbuntuFromBaseImageJob(configName, baseImageTag, archsToTest) {
  return {
    stage("${configName}") {
      // Notify Gitlab about the build and tests it will be running
      // so it doesn't the build successful before it start running them
      // and we can also see issue if the build / test is never run.
      updateGitlabCommitStatus(name: "build-${configName}", state: "pending")
      for (arch in archsToTest.split(';')) {
        updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: "pending")
      }
      build job: "ubuntu_build_template_CI",
      parameters: [
        string(name: 'configName', value: "${configName}"),
        string(name: 'baseImageTag', value: "${baseImageTag}"),
        string(name: 'targetImageTag',
               value: "${wispDockerRegistryFull}:${branchRef}-${BUILD_ID}-${configName}"),
        string(name: 'archsToTest', value: "${archsToTest}"),
        string(name: 'sourceBranch', value: "${env.gitlabSourceBranch}"),
        string(name: 'repoUrl', value: "${scm.userRemoteConfigs[0].url}"),
        string(name: 'commitHash', value: "${commitHash}")
      ],
      // This node doesn't need to be held while builds and tests run.
      wait: false,
      // Success of this script depend only on successful launch,
      // Not successful builds and tests.
      propagate: false
    }
  }
}
